// SPDX-License-Identifier: Mulan PSL v2
/*
 * Copyright (c) 2025 Huawei Technologies Co., Ltd.
 * This software is licensed under Mulan PSL v2.
 * You can use this software according to the terms and conditions of the Mulan PSL v2.
 * You may obtain a copy of Mulan PSL v2 at:
 *         http://license.coscl.org.cn/MulanPSL2
 *
 * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND,
 * EITHER EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT,
 * MERCHANTABILITY OR FIT FOR A PARTICULAR PURPOSE.
 * See the Mulan PSL v2 for more details.
 */

use std::{env, fs, io, path::PathBuf, process};

use anyhow::{Context, Result};
use once_cell::sync::OnceCell;
use parking_lot::Mutex;
use tracing::{level_filters::LevelFilter, Level};
use tracing_appender::non_blocking::WorkerGuard;
use tracing_subscriber::{fmt, prelude::*, EnvFilter, Layer};

const MODULE_LEVELS: &[(&str, Level)] = &[("xgpu_common::ipc", Level::ERROR)];

static GLOBAL_LOG_GUARDS: OnceCell<Mutex<Option<Vec<WorkerGuard>>>> = OnceCell::new();

/// A struct that implements `std::io::Write` but delays file creation until the first write.
#[derive(Debug)]
struct LazyFileWriter {
    path: PathBuf,
    file: OnceCell<fs::File>,
}

impl LazyFileWriter {
    fn new(path: PathBuf) -> Self {
        Self {
            path,
            file: OnceCell::new(),
        }
    }
}

impl io::Write for LazyFileWriter {
    fn write(&mut self, buf: &[u8]) -> io::Result<usize> {
        self.file
            .get_or_try_init(|| {
                if let Some(parent) = self.path.parent() {
                    fs::create_dir_all(parent)?;
                }
                fs::File::create(&self.path)
            })?
            .write(buf)
    }

    fn flush(&mut self) -> io::Result<()> {
        if let Some(file) = self.file.get_mut() {
            file.flush()?;
        }
        Ok(())
    }
}

/// Initializes the logging system.
pub fn initialize() -> Result<()> {
    if GLOBAL_LOG_GUARDS.set(Mutex::new(Some(Vec::new()))).is_err() {
        anyhow::bail!("Logger has already been initialized");
    }

    // --- 1. Create file logging layer ---
    let log_path = {
        let log_dir = env::current_dir().context("Failed to get current working directory")?;

        let exec_path = env::current_exe().context("Failed to get current executable path")?;
        let exec_name = exec_path
            .file_stem()
            .map(|name| name.to_string_lossy())
            .context("Failed to get binary name from executable path")?;
        let pid = process::id();

        log_dir.join(format!("{exec_name}_{pid}.log"))
    };

    let (file_writer, file_guard) = tracing_appender::non_blocking(LazyFileWriter::new(log_path));

    let env_filter = MODULE_LEVELS.iter().try_fold(
        EnvFilter::from_default_env(),
        |filter, (module_path, level)| {
            let directive_str = format!("{module_path}={level}");
            directive_str
                .parse()
                .with_context(|| {
                    format!(
                        "Failed to parse built-in log directive: '{}'",
                        directive_str
                    )
                })
                .map(|directive| filter.add_directive(directive))
        },
    )?;

    let file_layer = fmt::layer()
        .with_writer(file_writer)
        .with_ansi(false)
        .with_filter(env_filter);

    // --- 2. Create stderr logging layer ---
    let (stderr_writer, stderr_guard) = tracing_appender::non_blocking(io::stderr());
    let stderr_layer = fmt::layer()
        .with_writer(stderr_writer)
        .with_ansi(true)
        .with_filter(LevelFilter::from_level(Level::WARN));

    // --- 3. Initialize the global subscriber ---
    tracing_subscriber::registry()
        .with(file_layer)
        .with(stderr_layer)
        .init();

    // --- 4. Store ALL guards for unified, graceful shutdown ---
    let mut guard_lock = GLOBAL_LOG_GUARDS
        .get()
        .context("Logger guards were not initialized correctly")?
        .lock();

    let guards_vec = guard_lock
        .as_mut()
        .context("Logger guards were unexpectedly taken during initialization")?;

    guards_vec.push(file_guard);
    guards_vec.push(stderr_guard);

    Ok(())
}
